"""
Steady state detection and stability analysis

This module provides tools for detecting when a PDE solution has reached
a steady state and analyzing the stability of equilibrium solutions.
"""

using ..Fields
using ..PDEs
using LinearAlgebra
using Statistics

"""
    SteadyStateTracker

Tracks evolution towards steady state and detects convergence.
"""
struct SteadyStateTracker
    tolerance::Float64
    window_size::Int
    check_frequency::Int
    
    # State tracking
    field_history::Vector{Any}  # Store recent states for comparison
    time_history::Vector{Float64}
    norm_history::Vector{Float64}
    change_rate_history::Vector{Float64}
    
    # Convergence detection
    is_converged::Bool
    convergence_time::Float64
    steady_state_value::Any
    
    # Statistics
    current_step::Int
    
    function SteadyStateTracker(; tolerance=1e-6, window_size=50, check_frequency=10)
        new(tolerance, window_size, check_frequency,
            Any[], Float64[], Float64[], Float64[],
            false, -1.0, nothing, 0)
    end
end

"""
    track!(tracker::SteadyStateTracker, pde::AbstractPDE, state, t)

Monitor approach to steady state.
"""
function track!(tracker::SteadyStateTracker, pde::AbstractPDE, state, t)
    tracker.current_step += 1
    
    if tracker.current_step % tracker.check_frequency == 0
        # Store current state
        push!(tracker.field_history, deepcopy(state))
        push!(tracker.time_history, t)
        
        # Compute state norm
        state_norm = compute_state_norm(state)
        push!(tracker.norm_history, state_norm)
        
        # Keep only recent history
        if length(tracker.field_history) > tracker.window_size
            popfirst!(tracker.field_history)
            popfirst!(tracker.time_history)
            popfirst!(tracker.norm_history)
        end
        
        # Compute change rate
        if length(tracker.field_history) >= 2
            recent_states = tracker.field_history[end-1:end]
            recent_times = tracker.time_history[end-1:end]
            
            state_change = compute_state_difference(recent_states[2], recent_states[1])
            time_change = recent_times[2] - recent_times[1]
            change_rate = state_change / time_change
            
            push!(tracker.change_rate_history, change_rate)
            
            # Keep change rate history manageable
            if length(tracker.change_rate_history) > tracker.window_size
                popfirst!(tracker.change_rate_history)
            end
        end
        
        # Check for convergence
        if !tracker.is_converged && length(tracker.change_rate_history) >= 10
            if check_convergence(tracker)
                tracker.is_converged = true
                tracker.convergence_time = t
                tracker.steady_state_value = deepcopy(state)
                @info "Steady state detected at t = $t"
            end
        end
    end
end

"""
    compute_state_norm(state)

Compute an appropriate norm of the state.
"""
function compute_state_norm(state)
    if isa(state, ScalarField)
        return sqrt(sum(data(state).^2))
    elseif isa(state, FieldCollection)
        total_norm_sq = 0.0
        for field in state
            total_norm_sq += sum(data(field).^2)
        end
        return sqrt(total_norm_sq)
    else
        return norm(state)
    end
end

"""
    compute_state_difference(state1, state2)

Compute difference between two states.
"""
function compute_state_difference(state1, state2)
    if isa(state1, ScalarField) && isa(state2, ScalarField)
        diff = data(state1) - data(state2)
        return sqrt(sum(diff.^2))
    elseif isa(state1, FieldCollection) && isa(state2, FieldCollection)
        total_diff_sq = 0.0
        for (f1, f2) in zip(state1, state2)
            diff = data(f1) - data(f2)
            total_diff_sq += sum(diff.^2)
        end
        return sqrt(total_diff_sq)
    else
        return norm(state1 - state2)
    end
end

"""
    check_convergence(tracker::SteadyStateTracker)

Check if the system has converged to steady state.
"""
function check_convergence(tracker::SteadyStateTracker)
    # Need sufficient history
    if length(tracker.change_rate_history) < 10
        return false
    end
    
    # Check if change rates are consistently small
    recent_rates = tracker.change_rate_history[end-9:end]
    max_rate = maximum(recent_rates)
    avg_rate = mean(recent_rates)
    
    # Convergence criteria
    return max_rate < tracker.tolerance && avg_rate < tracker.tolerance / 2
end

"""
    get_convergence_info(tracker::SteadyStateTracker)

Get information about convergence status.
"""
function get_convergence_info(tracker::SteadyStateTracker)
    return (
        converged = tracker.is_converged,
        convergence_time = tracker.convergence_time,
        current_change_rate = isempty(tracker.change_rate_history) ? 
                             Inf : tracker.change_rate_history[end],
        tolerance = tracker.tolerance
    )
end

"""
    StabilityAnalyzer

Analyze linear stability of steady states.
"""
struct StabilityAnalyzer
    linearization_method::Symbol  # :finite_difference, :symbolic, :automatic
    perturbation_size::Float64
    eigenvalue_tolerance::Float64
    
    function StabilityAnalyzer(; method=:finite_difference, perturbation=1e-8, 
                              eigenvalue_tol=1e-10)
        new(method, perturbation, eigenvalue_tol)
    end
end

"""
    analyze_stability(analyzer::StabilityAnalyzer, pde::AbstractPDE, steady_state)

Perform linear stability analysis of a steady state.
"""
function analyze_stability(analyzer::StabilityAnalyzer, pde::AbstractPDE, steady_state)
    if analyzer.linearization_method == :finite_difference
        return finite_difference_stability(analyzer, pde, steady_state)
    elseif analyzer.linearization_method == :symbolic
        return symbolic_stability(analyzer, pde, steady_state)
    else
        error("Stability analysis method $(analyzer.linearization_method) not implemented")
    end
end

"""
    finite_difference_stability(analyzer, pde, steady_state)

Compute linearization using finite differences.
"""
function finite_difference_stability(analyzer::StabilityAnalyzer, pde::AbstractPDE, 
                                    steady_state)
    # Get steady state evolution rate (should be ~0)
    f0 = evolution_rate(pde, steady_state, 0.0)
    
    # Compute Jacobian via finite differences
    jacobian = compute_jacobian_fd(pde, steady_state, analyzer.perturbation_size)
    
    # Compute eigenvalues
    eigenvals = eigvals(jacobian)
    
    # Classify stability
    max_real_part = maximum(real.(eigenvals))
    stability = classify_stability(max_real_part, analyzer.eigenvalue_tolerance)
    
    return (
        eigenvalues = eigenvals,
        max_real_eigenvalue = max_real_part,
        stability = stability,
        jacobian = jacobian
    )
end

"""
    compute_jacobian_fd(pde, steady_state, h)

Compute Jacobian matrix using finite differences.
"""
function compute_jacobian_fd(pde::AbstractPDE, steady_state, h::Float64)
    # This is a simplified version - full implementation would need
    # careful handling of field structures
    
    if isa(steady_state, ScalarField)
        data_vec = vec(data(steady_state))
        n = length(data_vec)
        jacobian = zeros(n, n)
        
        for i in 1:n
            # Perturb i-th component
            perturbed_data = copy(data_vec)
            perturbed_data[i] += h
            
            # Create perturbed state
            perturbed_state = ScalarField(reshape(perturbed_data, size(data(steady_state))), 
                                        grid(steady_state))
            
            # Compute evolution rate
            f_plus = evolution_rate(pde, perturbed_state, 0.0)
            f_minus = evolution_rate(pde, steady_state, 0.0)
            
            # Finite difference approximation
            jacobian[:, i] = vec(data(f_plus) - data(f_minus)) / h
        end
        
        return jacobian
    else
        error("Jacobian computation for FieldCollection not yet implemented")
    end
end

"""
    classify_stability(max_real_eigenvalue, tolerance)

Classify stability based on largest real eigenvalue.
"""
function classify_stability(max_real_eigenvalue::Float64, tolerance::Float64)
    if max_real_eigenvalue < -tolerance
        return :stable
    elseif max_real_eigenvalue > tolerance
        return :unstable
    else
        return :marginal
    end
end

"""
    PeriodicOrbitDetector

Detect periodic orbits and limit cycles.
"""
struct PeriodicOrbitDetector
    min_period::Float64
    max_period::Float64
    tolerance::Float64
    fourier_threshold::Float64
    
    # State tracking
    time_series::Vector{Float64}
    value_series::Vector{Float64}
    frequencies::Vector{Float64}
    amplitudes::Vector{Float64}
    
    function PeriodicOrbitDetector(; min_period=1.0, max_period=100.0, 
                                  tolerance=1e-6, fourier_threshold=0.1)
        new(min_period, max_period, tolerance, fourier_threshold,
            Float64[], Float64[], Float64[], Float64[])
    end
end

"""
    track!(detector::PeriodicOrbitDetector, pde::AbstractPDE, state, t)

Track time series for periodic orbit detection.
"""
function track!(detector::PeriodicOrbitDetector, pde::AbstractPDE, state, t)
    # Store time and a scalar measure of the state
    push!(detector.time_series, t)
    
    # Use L2 norm as scalar measure
    scalar_value = compute_state_norm(state)
    push!(detector.value_series, scalar_value)
    
    # Keep only recent history for efficiency
    max_length = 10000
    if length(detector.time_series) > max_length
        detector.time_series = detector.time_series[end-max_length+1:end]
        detector.value_series = detector.value_series[end-max_length+1:end]
    end
end

"""
    detect_periodicity(detector::PeriodicOrbitDetector)

Analyze time series for periodicity using Fourier analysis.
"""
function detect_periodicity(detector::PeriodicOrbitDetector)
    if length(detector.value_series) < 100
        return (periodic = false, period = 0.0, confidence = 0.0)
    end
    
    # Simplified FFT analysis
    # In practice, would use FFTW.jl for efficient computation
    
    # For now, return simple autocorrelation-based detection
    return autocorrelation_periodicity(detector)
end

"""
    autocorrelation_periodicity(detector::PeriodicOrbitDetector)

Detect periodicity using autocorrelation.
"""
function autocorrelation_periodicity(detector::PeriodicOrbitDetector)
    values = detector.value_series
    times = detector.time_series
    
    if length(values) < 50
        return (periodic = false, period = 0.0, confidence = 0.0)
    end
    
    # Compute autocorrelation at different lags
    dt = times[2] - times[1]  # Assume uniform sampling
    max_lag = min(length(values) ÷ 2, Int(detector.max_period / dt))
    
    autocorr = zeros(max_lag)
    mean_val = mean(values)
    var_val = var(values)
    
    for lag in 1:max_lag
        correlation = 0.0
        count = 0
        for i in 1:(length(values) - lag)
            correlation += (values[i] - mean_val) * (values[i + lag] - mean_val)
            count += 1
        end
        autocorr[lag] = correlation / (count * var_val)
    end
    
    # Find peaks in autocorrelation
    peak_threshold = 0.5
    peaks = findall(i -> i > 1 && i < length(autocorr) && 
                         autocorr[i] > autocorr[i-1] && 
                         autocorr[i] > autocorr[i+1] && 
                         autocorr[i] > peak_threshold, 
                    2:(length(autocorr)-1))
    
    if isempty(peaks)
        return (periodic = false, period = 0.0, confidence = 0.0)
    end
    
    # Most likely period is the first significant peak
    first_peak = peaks[1]
    period = first_peak * dt
    confidence = autocorr[first_peak]
    
    return (periodic = confidence > peak_threshold, period = period, confidence = confidence)
end

"""
    BifurcationAnalyzer

Analyze bifurcations in parameter space.
"""
struct BifurcationAnalyzer
    parameter_name::Symbol
    parameter_range::Vector{Float64}
    steady_states::Vector{Any}
    eigenvalues::Vector{Vector{ComplexF64}}
    bifurcation_points::Vector{Float64}
    
    function BifurcationAnalyzer(param_name::Symbol, param_range::Vector{Float64})
        new(param_name, param_range, Any[], Vector{ComplexF64}[], Float64[])
    end
end

"""
    analyze_bifurcation(analyzer::BifurcationAnalyzer, pde_template::Function, 
                       initial_state, solver_config)

Perform bifurcation analysis by varying a parameter.
"""
function analyze_bifurcation(analyzer::BifurcationAnalyzer, pde_template::Function,
                            initial_state, solver_config=nothing)
    stability_analyzer = StabilityAnalyzer()
    
    for param_value in analyzer.parameter_range
        # Create PDE with current parameter value
        pde = pde_template(param_value)
        
        # Find steady state (simplified - would need actual steady state solver)
        steady_state = find_steady_state(pde, initial_state, solver_config)
        push!(analyzer.steady_states, steady_state)
        
        # Analyze stability
        stability_result = analyze_stability(stability_analyzer, pde, steady_state)
        push!(analyzer.eigenvalues, stability_result.eigenvalues)
        
        # Check for bifurcation (eigenvalue crossing imaginary axis)
        if length(analyzer.eigenvalues) > 1
            prev_max_real = maximum(real.(analyzer.eigenvalues[end-1]))
            curr_max_real = maximum(real.(analyzer.eigenvalues[end]))
            
            if prev_max_real < 0 && curr_max_real > 0
                # Bifurcation detected
                push!(analyzer.bifurcation_points, param_value)
                @info "Bifurcation detected at $(analyzer.parameter_name) = $param_value"
            end
        end
    end
    
    return analyzer
end

"""
    find_steady_state(pde, initial_state, solver_config)

Find steady state of PDE (simplified implementation).
"""
function find_steady_state(pde::AbstractPDE, initial_state, solver_config)
    # This is a placeholder - would need actual steady state finding algorithm
    # Could use Newton's method on F(u) = 0 where F is the evolution rate
    return initial_state  # Simplified
end

# Export steady state and stability analysis tools
export SteadyStateTracker, StabilityAnalyzer, PeriodicOrbitDetector, BifurcationAnalyzer
export analyze_stability, detect_periodicity, analyze_bifurcation
export get_convergence_info, classify_stability